use std::mem;
use windows::Win32::System::LibraryLoader::{GetProcAddress, LoadLibraryW};
use windows::core::HSTRING;

use crate::classes::*;
use crate::enums::*;
use crate::registry_manager::*;

pub struct SetProcessMitigationsCommand;

impl SetProcessMitigationsCommand {
    const BLOCK_IN_SYSTEM_MODE: &'static [MitigationOptions] = &[
        MitigationOptions::EnableExportAddressFilter,
        MitigationOptions::AuditEnableExportAddressFilter,
        MitigationOptions::EnableExportAddressFilterPlus,
        MitigationOptions::AuditEnableExportAddressFilterPlus,
        MitigationOptions::EnableImportAddressFilter,
        MitigationOptions::AuditEnableImportAddressFilter,
        MitigationOptions::EnableRopStackPivot,
        MitigationOptions::AuditEnableRopStackPivot,
        MitigationOptions::EnableRopCallerCheck,
        MitigationOptions::AuditEnableRopCallerCheck,
        MitigationOptions::EnableRopSimExec,
        MitigationOptions::AuditEnableRopSimExec,
        MitigationOptions::DisallowChildProcessCreation,
        MitigationOptions::AuditChildProcess,
    ];

    pub fn import_mitigation(xml_path: &str) -> i32 {
        unsafe {
            if let Ok(mitigation_config_dll) =
                LoadLibraryW(&HSTRING::from("MitigationConfiguration.dll"))
            {
                if let Some(proc_addr) =
                    GetProcAddress(mitigation_config_dll, windows::core::s!("ImportMitigation"))
                {
                    let func: extern "system" fn(*const u16) -> i32 = mem::transmute(proc_addr);
                    let wide_path: Vec<u16> =
                        xml_path.encode_utf16().chain(std::iter::once(0)).collect();
                    func(wide_path.as_ptr())
                } else {
                    -1
                }
            } else {
                -1
            }
        }
    }

    pub fn validate_xml_from_managed(file_name: &str, result: &mut bool) -> i32 {
        unsafe {
            if let Ok(mitigation_config_dll) =
                LoadLibraryW(&HSTRING::from("MitigationConfiguration.dll"))
            {
                if let Some(proc_addr) = GetProcAddress(
                    mitigation_config_dll,
                    windows::core::s!("ValidateXMLFromManaged"),
                ) {
                    let func: extern "system" fn(*const u16, *mut bool) -> i32 =
                        mem::transmute(proc_addr);
                    let wide_path: Vec<u16> =
                        file_name.encode_utf16().chain(std::iter::once(0)).collect();
                    func(wide_path.as_ptr(), result as *mut bool)
                } else {
                    -1
                }
            } else {
                -1
            }
        }
    }

    fn is_option_in_list(list: Option<&[MitigationOptions]>, option: MitigationOptions) -> bool {
        if let Some(list) = list {
            list.contains(&option)
        } else {
            false
        }
    }

    pub fn set_enable_and_disable(
        app_mitigations: &mut AppMitigations,
        _registry_manager: &mut RegistryManager,
        disable_list: Option<&[MitigationOptions]>,
        enable_list: Option<&[MitigationOptions]>,
        eaf_modules_list: Option<&[String]>,
        is_force: Option<&str>,
        is_remove: bool,
        is_reset: bool,
        is_system_mode: bool,
    ) -> Result<(), String> {
        if enable_list.is_some() || disable_list.is_some() {
            let policy_pairs = PolicyPairList::get_policy_pairs();

            for pair in policy_pairs.iter() {
                if !is_system_mode
                    || !Self::is_option_in_list(
                        Some(Self::BLOCK_IN_SYSTEM_MODE),
                        pair.parameter_name,
                    )
                {
                    if Self::is_option_in_list(enable_list, pair.parameter_name)
                        || Self::is_option_in_list(disable_list, pair.parameter_name)
                    {
                        let mut option_value = OptionValue::NotSet;

                        if Self::is_option_in_list(enable_list, pair.parameter_name) {
                            option_value = OptionValue::On;
                        } else if Self::is_option_in_list(disable_list, pair.parameter_name) {
                            option_value = OptionValue::Off;
                        }

                        if let Some(policy) = app_mitigations.policies.get_mut(&pair.policy_name) {
                            // Set the policy option value
                            policy.set_policy(pair.policy_option_name, option_value);

                            // Handle flags based on policy type
                            match policy {
                                AppPolicyContainer::Dep(p) => {
                                    if is_reset {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagReset as u32,
                                        );
                                    }
                                    if is_remove {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagRemove as u32,
                                        );
                                    }
                                    if let Some(force_str) = is_force {
                                        if force_str != "notset" {
                                            let force_value = force_str.to_lowercase() == "on";
                                            p.set_override_dep(force_value);
                                        }
                                    }
                                }
                                AppPolicyContainer::Aslr(p) => {
                                    if is_reset {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagReset as u32,
                                        );
                                    }
                                    if is_remove {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagRemove as u32,
                                        );
                                    }
                                    if let Some(force_str) = is_force {
                                        if force_str != "notset" {
                                            let force_value = force_str.to_lowercase() == "on";
                                            match pair.force_name {
                                                ForceName::ForceRelocateImages => p
                                                    .set_override_force_relocate_images(
                                                        force_value,
                                                    ),
                                                ForceName::BottomUp => {
                                                    p.set_override_bottom_up(force_value)
                                                }
                                                ForceName::HighEntropy => {
                                                    p.set_override_high_entropy(force_value)
                                                }
                                                _ => {}
                                            }
                                        }
                                    }
                                }
                                AppPolicyContainer::StrictHandle(p) => {
                                    if is_reset {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagReset as u32,
                                        );
                                    }
                                    if is_remove {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagRemove as u32,
                                        );
                                    }
                                    if let Some(force_str) = is_force {
                                        if force_str != "notset" {
                                            let force_value = force_str.to_lowercase() == "on";
                                            p.set_override_strict_handle(force_value);
                                        }
                                    }
                                }
                                AppPolicyContainer::SystemCalls(p) => {
                                    if is_reset {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagReset as u32,
                                        );
                                    }
                                    if is_remove {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagRemove as u32,
                                        );
                                    }
                                    if let Some(force_str) = is_force {
                                        if force_str != "notset" {
                                            let force_value = force_str.to_lowercase() == "on";
                                            match pair.force_name {
                                                ForceName::SystemCall => {
                                                    p.set_override_system_call(force_value)
                                                }
                                                ForceName::DisableFsctlSystem => {
                                                    p.set_override_fsctl_system_call(force_value)
                                                }
                                                _ => {}
                                            }
                                        }
                                    }
                                }
                                AppPolicyContainer::ExtensionPoints(p) => {
                                    if is_reset {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagReset as u32,
                                        );
                                    }
                                    if is_remove {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagRemove as u32,
                                        );
                                    }
                                    if let Some(force_str) = is_force {
                                        if force_str != "notset" {
                                            let force_value = force_str.to_lowercase() == "on";
                                            p.set_override_extension_point(force_value);
                                        }
                                    }
                                }
                                AppPolicyContainer::DynamicCode(p) => {
                                    if is_reset {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagReset as u32,
                                        );
                                    }
                                    if is_remove {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagRemove as u32,
                                        );
                                    }
                                    if let Some(force_str) = is_force {
                                        if force_str != "notset" {
                                            let force_value = force_str.to_lowercase() == "on";
                                            p.set_override_dynamic_code(force_value);
                                        }
                                    }
                                }
                                AppPolicyContainer::ControlFlowGuard(p) => {
                                    if is_reset {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagReset as u32,
                                        );
                                    }
                                    if is_remove {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagRemove as u32,
                                        );
                                    }
                                    if let Some(force_str) = is_force {
                                        if force_str != "notset" {
                                            let force_value = force_str.to_lowercase() == "on";
                                            match pair.force_name {
                                                ForceName::CFG => p.set_override_cfg(force_value),
                                                ForceName::StrictCFG => {
                                                    p.set_override_strict_cfg(force_value)
                                                }
                                                _ => {}
                                            }
                                        }
                                    }
                                }
                                AppPolicyContainer::SignedBinaries(p) => {
                                    if is_reset {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagReset as u32,
                                        );
                                    }
                                    if is_remove {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagRemove as u32,
                                        );
                                    }
                                    if let Some(force_str) = is_force {
                                        if force_str != "notset" {
                                            let force_value = force_str.to_lowercase() == "on";
                                            match pair.force_name {
                                                ForceName::MicrosoftSignedOnly => p
                                                    .set_override_microsoft_signed_only(
                                                        force_value,
                                                    ),
                                                ForceName::EnforceModuleDependencySigning => p
                                                    .set_override_enforce_module_dependency_signing(
                                                        force_value,
                                                    ),
                                                _ => {}
                                            }
                                        }
                                    }
                                }
                                AppPolicyContainer::Fonts(p) => {
                                    if is_reset {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagReset as u32,
                                        );
                                    }
                                    if is_remove {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagRemove as u32,
                                        );
                                    }
                                    if let Some(force_str) = is_force {
                                        if force_str != "notset" {
                                            let force_value = force_str.to_lowercase() == "on";
                                            p.set_override_font_disable(force_value);
                                        }
                                    }
                                }
                                AppPolicyContainer::ImageLoad(p) => {
                                    if is_reset {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagReset as u32,
                                        );
                                    }
                                    if is_remove {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagRemove as u32,
                                        );
                                    }
                                    if let Some(force_str) = is_force {
                                        if force_str != "notset" {
                                            let force_value = force_str.to_lowercase() == "on";
                                            match pair.force_name {
                                                ForceName::BlockRemoteImageLoads => p
                                                    .set_override_block_remote_image_loads(
                                                        force_value,
                                                    ),
                                                ForceName::BlockLowLabel => {
                                                    p.set_override_block_low_label(force_value)
                                                }
                                                ForceName::PreferSystem32 => {
                                                    p.set_override_prefer_system32(force_value)
                                                }
                                                _ => {}
                                            }
                                        }
                                    }
                                }
                                AppPolicyContainer::Payload(p) => {
                                    if is_reset {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagReset as u32,
                                        );
                                    }
                                    if is_remove {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagRemove as u32,
                                        );
                                    }
                                    if let Some(force_str) = is_force {
                                        if force_str != "notset" {
                                            let force_value = force_str.to_lowercase() == "on";
                                            match pair.force_name {
                                                ForceName::EnableExportAddressFilter => p
                                                    .set_override_enable_export_address_filter(
                                                        force_value,
                                                    ),
                                                ForceName::EnableExportAddressFilterPlus => p
                                                    .set_override_enable_export_address_filter_plus(
                                                        force_value,
                                                    ),
                                                ForceName::EnableImportAddressFilter => p
                                                    .set_override_enable_import_address_filter(
                                                        force_value,
                                                    ),
                                                ForceName::EnableRopStackPivot => p
                                                    .set_override_enable_rop_stack_pivot(
                                                        force_value,
                                                    ),
                                                ForceName::EnableRopCallerCheck => p
                                                    .set_override_enable_rop_caller_check(
                                                        force_value,
                                                    ),
                                                ForceName::EnableRopSimExec => {
                                                    p.set_override_enable_rop_sim_exec(force_value)
                                                }
                                                _ => {}
                                            }
                                        }
                                    }
                                }
                                AppPolicyContainer::ChildProcess(p) => {
                                    if is_reset {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagReset as u32,
                                        );
                                    }
                                    if is_remove {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagRemove as u32,
                                        );
                                    }
                                    if let Some(force_str) = is_force {
                                        if force_str != "notset" {
                                            let force_value = force_str.to_lowercase() == "on";
                                            p.set_override_child_process(force_value);
                                        }
                                    }
                                }
                                AppPolicyContainer::UserShadowStack(p) => {
                                    if is_reset {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagReset as u32,
                                        );
                                    }
                                    if is_remove {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagRemove as u32,
                                        );
                                    }
                                    if let Some(force_str) = is_force {
                                        if force_str != "notset" {
                                            let force_value = force_str.to_lowercase() == "on";
                                            p.set_override_user_shadow_stack(force_value);
                                        }
                                    }
                                }
                                AppPolicyContainer::SEHOP(p) => {
                                    if is_reset {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagReset as u32,
                                        );
                                    }
                                    if is_remove {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagRemove as u32,
                                        );
                                    }
                                    if let Some(force_str) = is_force {
                                        if force_str != "notset" {
                                            let force_value = force_str.to_lowercase() == "on";
                                            p.set_override_sehop(force_value);
                                        }
                                    }
                                }
                                AppPolicyContainer::Heap(p) => {
                                    if is_reset {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagReset as u32,
                                        );
                                    }
                                    if is_remove {
                                        p.set_flags(
                                            p.flags() | ProcessMitigationFlags::FlagRemove as u32,
                                        );
                                    }
                                    if let Some(force_str) = is_force {
                                        if force_str != "notset" {
                                            let force_value = force_str.to_lowercase() == "on";
                                            p.set_override_heap(force_value);
                                        }
                                    }
                                }
                            }
                        }
                    }
                } else if Self::is_option_in_list(enable_list, pair.parameter_name)
                    || Self::is_option_in_list(disable_list, pair.parameter_name)
                {
                    println!(
                        "{} is not supported in system mode.",
                        pair.parameter_name.to_parameter_name()
                    );
                }
            }
        }

        if let Some(eaf_modules) = eaf_modules_list {
            // Handle EAF modules
            if let Some(AppPolicyContainer::Payload(payload_policy)) =
                app_mitigations.policies.get_mut(&PolicyName::Payload)
            {
                payload_policy.eaf_modules.clear();
                for module in eaf_modules {
                    if !module.is_empty() && !payload_policy.eaf_modules.contains(module) {
                        payload_policy.eaf_modules.push(module.clone());
                    }
                }

                if payload_policy.enable_export_address_filter_plus() != OptionValue::On {
                    println!("EAF+ modules specified, but EAF+ is not enabled");
                }
            }
        }

        if disable_list.is_some() || enable_list.is_some() || eaf_modules_list.is_some() {
            RegistryManager::write_policy_to_registry(app_mitigations)
        } else {
            Ok(())
        }
    }
}
